package net.nohaven.proj.javeau.crypt.streams;

import java.io.IOException;
import java.io.OutputStream;

import net.nohaven.proj.javeau.crypt.engine.Phelix;

public final class PhelixOutputStream extends OutputStream {
	private static final int BUFFER_DIM = 4096;

	private OutputStream out;

	private Phelix encoder = new Phelix(Phelix.PHELIX_MAC_SIZE);

	private long size;

	private byte[] buffer = new byte[BUFFER_DIM];
	private byte[] bufferOut = new byte[BUFFER_DIM];

	private int idx;

	private boolean closeUnderlying;
	private boolean closed = false;

	public PhelixOutputStream(OutputStream outs, byte[] key, byte[] iv)
			throws IOException {
		this(outs, key, iv, true);
	}

	public PhelixOutputStream(OutputStream outs, byte[] key, byte[] iv,
			boolean closeUnderlying) throws IOException {
		this.out = outs;
		this.closeUnderlying = closeUnderlying;

		byte[] nonce = new byte[16];
		System.arraycopy(iv, 0, nonce, 0, Math.min(iv.length, 16));

		encoder.setupKey(Phelix.MODE_ENCRYPT, key);
		encoder.setupNonce(nonce);

		size = 0;
	}

	private void flushBuffer() throws IOException {
		encoder.process(buffer, 0, bufferOut, 0, idx);

		out.write(bufferOut, 0, idx);
		size += idx;
		idx = 0;
	}

	public void write(byte[] bytes, int off, int len) throws IOException {
		while (true) {
			if (len == 0)
				break;
			int rem = BUFFER_DIM - idx;
			if (rem > len) {
				System.arraycopy(bytes, off, buffer, idx, len);
				idx += len;
				break;
			}
			System.arraycopy(bytes, off, buffer, idx, rem);
			idx = BUFFER_DIM;
			flushBuffer();
			len -= rem;
			off += rem;
		}
	}

	public void write(byte[] bytes) throws IOException {
		write(bytes, 0, bytes.length);
	}

	public void write(int bbyte) throws IOException {
		buffer[idx++] = (byte) bbyte;
		if (idx == buffer.length)
			flushBuffer();
	}

	public void close() throws IOException {
		if (isClosed())
			return;
		closed = true;
		flushBuffer();
		if (closeUnderlying)
			out.close();
	}

	public void flush() throws IOException {
		out.flush();
	}

	public long getSize() {
		return size;
	}

	public boolean isClosed() {
		return closed;
	}

	public byte[] getMac() {
		byte[] mac = new byte[Phelix.PHELIX_MAC_SIZE >> 3];
		encoder.finalize(mac, 0);
		return mac;
	}

	//	public static void main(String[] args) throws Exception {
	//		byte[] key = new byte[32];
	//		byte[] iv = new byte[16];
	//
	//		byte[] data = new byte[224205];
	//		byte[] data2 = new byte[data.length];
	//		byte[] mac, mac2;
	//
	//		new Random().nextBytes(key);
	//		new Random().nextBytes(iv);
	//		new Random().nextBytes(data);
	//
	//		ByteArrayOutputStream baos = new ByteArrayOutputStream();
	//
	//		PhelixOutputStream cos = new PhelixOutputStream(baos, key, iv, false);
	//		cos.write(data);
	//		cos.flush();
	//		cos.close();
	//		mac = cos.getMac();
	//
	//		ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
	//
	//		PhelixInputStream cis = new PhelixInputStream(bais, key, iv, cos
	//				.getSize(), false);
	//		cis.read(data2);
	//		cis.close();
	//		mac2 = cis.getMac();
	//
	//		System.out.println(Arrays.equals(data, data2));
	//		System.out.println(Arrays.equals(mac, mac2));
	//	}
}
